__all__ = ["tb"]

import os
import torch.utils.tensorboard as _tensorboard

class tb:
    def __enter__(self):
        if not os.path.exists("logs"):
            os.mkdir("logs")
        else:
            for file in os.listdir("logs"):
                os.remove(os.path.join("logs", file))

        self.writer = _tensorboard.SummaryWriter("logs")
        return self.writer

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is None:
            self.writer.close()
            try:
                os.system("tensorboard --logdir logs --port 6006")
            except KeyboardInterrupt:
                exit(0)